Numerical Quantum Mechanics

Time-Independent Schrodinger Equation using Finite Difference

Fadjar Fathurrahman

Mariya Al Qibtiya Nasution

In [12]:
using Printf
In [1]:
using LinearAlgebra
In [3]:
using SparseArrays

For calculating several lowest eigensolutions of large sparse matrix (and other iterative solvers).

In [4]:
using IterativeSolvers

This function emulates speye function which is also found in MATLAB.

In [5]:
speye(N::Int64) = sparse(Matrix(1.0I, N, N));

For plotting (need LaTeX and pgfplots to be installed previously)

In [2]:
using PGFPlotsX

Introduction

Let's consider again harmonic potential, but now it is in 2d: $$ V(x,y) = \frac{1}{2} m \omega^2 \left( x^2 + y^2 \right) $$

We will use Hartree atomic units: $\hbar = 1$, $m = 1$ and we take $\omega = 1$.

We will compare our numerical eigenvalues with nalytical eigenvalues:

$$ E_{n_{x},n_{y}} = \hbar^2 \omega^2 \left( n_{x} + n_{y} + 1 \right) $$

with $n_x$ and $n_y$ start from $0, 1, 2, 3, \ldots$.

Grid initialization

Let's first initialize the grid in $x$-direction:

In [6]:
Nx = 5
x_min = -5.0
x_max = 5.0
Lx = x_max - x_min
hx = Lx/(Nx-1)
x = zeros(Float64,Nx)
for i in 1:Nx
    x[i] = x_min + (i-1)*hx
end
In [9]:
x
Out[9]:
5-element Array{Float64,1}:
 -5.0
 -2.5
  0.0
  2.5
  5.0

Similarly, we can initialize the grid in $y$-direction:

In [7]:
Ny = 3
y_min = -5.0
y_max = 5.0
Ly = y_max - y_min
hy = Ly/(Ny-1)
y = zeros(Float64,Ny)
for j in 1:Ny
    y[j] = y_min + (j-1)*hy
end
In [10]:
y
Out[10]:
3-element Array{Float64,1}:
 -5.0
  0.0
  5.0

Now, we will define additional variables to describe our two dimensional grid.

We are essentially build a mapping between 2 indices to 1 indices

In [8]:
Npoints = Nx*Ny
r = zeros(2,Npoints)
ip = 0
idx_ip2xy = zeros(Int64,2,Npoints)
idx_xy2ip = zeros(Int64,Nx,Ny)
for j in 1:Ny
    for i in 1:Nx
        ip = ip + 1
        r[1,ip] = x[i]
        r[2,ip] = y[j]
        idx_ip2xy[1,ip] = i
        idx_ip2xy[2,ip] = j
        idx_xy2ip[i,j] = ip
    end
end

The grid points are contained in the variable r:

In [16]:
for ip in 1:Npoints
    @printf("Point #%3d: [%10.3f,%10.3f]\n", ip, r[1,ip], r[2,ip])
end
Point #  1: [    -5.000,    -5.000]
Point #  2: [    -2.500,    -5.000]
Point #  3: [     0.000,    -5.000]
Point #  4: [     2.500,    -5.000]
Point #  5: [     5.000,    -5.000]
Point #  6: [    -5.000,     0.000]
Point #  7: [    -2.500,     0.000]
Point #  8: [     0.000,     0.000]
Point #  9: [     2.500,     0.000]
Point # 10: [     5.000,     0.000]
Point # 11: [    -5.000,     5.000]
Point # 12: [    -2.500,     5.000]
Point # 13: [     0.000,     5.000]
Point # 14: [     2.500,     5.000]
Point # 15: [     5.000,     5.000]

Visualization of harmonic potential

In [17]:
function pot_harmonic(x::Float64, y::Float64; ω=1.0)
    return 0.5 * ω^2 *( x^2 + y^2 )
end
Out[17]:
pot_harmonic (generic function with 1 method)

For visualization purpose, we will use a larger grid size:

In [36]:
Nx = 50
x_min = -5.0
x_max = 5.0
Lx = x_max - x_min
hx = Lx/(Nx-1)
x = zeros(Float64,Nx)
for i in 1:Nx
    x[i] = x_min + (i-1)*hx
end
#
Ny = 50
y_min = -5.0
y_max = 5.0
Ly = y_max - y_min
hy = Ly/(Ny-1)
y = zeros(Float64,Ny)
for j in 1:Ny
    y[j] = y_min + (j-1)*hy
end
#
Npoints = Nx*Ny
r = zeros(2,Npoints)
ip = 0
idx_ip2xy = zeros(Int64,2,Npoints)
idx_xy2ip = zeros(Int64,Nx,Ny)
for j in 1:Ny
    for i in 1:Nx
        ip = ip + 1
        r[1,ip] = x[i]
        r[2,ip] = y[j]
        idx_ip2xy[1,ip] = i
        idx_ip2xy[2,ip] = j
        idx_xy2ip[i,j] = ip
    end
end
#
Vpot = zeros(Float64,Npoints)
for ip in 1:Npoints
    Vpot[ip] = pot_harmonic(r[1,ip], r[2,ip])
end
In [37]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(x, y, reshape(Vpot, Nx, Ny) )
  )
)
Out[37]:

Visualization of Gaussian potential

Let's consider another potential, i.e. Gaussian potential in 2d, centered at $\left( x_c, y_c \right)$:

$$ V_{c}(x,y) = -A\mathrm{e}^{-\alpha\left[ (x-x_c)^2 + (y-y_c)^2 \right]} $$
In [29]:
function pot_gaussian( x::Float64, y::Float64; A=1.0, α=1.0, xc=0.0, yc=0.0)
    return -A*exp(-α*((x-xc)^2 + (y-yc)^2))
end
Out[29]:
pot_gaussian (generic function with 1 method)
In [46]:
x1 = 1.0
y1 = 1.0
x2 = -1.0
y2 = -1.0
Vpot_gaussian = zeros(Float64,Npoints)
for ip in 1:Npoints
    Vpot_gaussian[ip] = pot_gaussian(r[1,ip], r[2,ip], xc=x1, yc=y1, α=0.75) + 
                        pot_gaussian(r[1,ip], r[2,ip], xc=x2, yc=y2, α=0.75)
end
In [48]:
@pgf Axis({ height = "10cm", width = "15cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(x, y, reshape(Vpot_gaussian, Nx, Ny) )
  )
)
Out[48]:
In [ ]:

Laplacian operator in 2d

We will start our discussion with a representation of Laplacian operators in 2d, which will be built from one dimensional second derivative operators in one dimension.

We need representation of Laplacian operator $\nabla^2$ in two dimension. We can use Kronecker sum, which in turns can be represented using Kronecker products:

$$ \begin{aligned} L & = D^{(2)}_{x} \oplus D^{(2)}_{y} \\ & = D^{(2)}_{x} \otimes \mathbf{I}_{y} + \mathbf{I}_{x} \otimes D^{(2)}_{y} \end{aligned} $$

where $D^{(2)}_{x}$ and $D^{(2)}_{y}$ are representation of second differential operators in $x$ dan $y$ direction and $\mathbf{I}_{x}$ and $\mathbf{I}_{y}$ are identity operators (matrix) in $x$ and $y$ direction.

Let's define again a function to build second differential matrix using 3 point finite difference approximation:

In [50]:
function build_D2_matrix(N::Int64, h::Float64)
    mat = zeros(Float64,N,N)
    for i = 1:N
        mat[i,i] = -2.0
        if i != N
            mat[i,i+1] = 1.0
            mat[i+1,i] = mat[i,i+1]
        end
    end
    return mat/h^2
end
Out[50]:
build_D2_matrix (generic function with 1 method)

Let's try first with a small grid size:

In [53]:
Nx = 4
x_min = -5.0
x_max = 5.0
Lx = x_max - x_min
hx = Lx/(Nx-1)
x = zeros(Float64,Nx)
for i in 1:Nx
    x[i] = x_min + (i-1)*hx
end
#
Ny = 5
y_min = -5.0
y_max = 5.0
Ly = y_max - y_min
hy = Ly/(Ny-1)
y = zeros(Float64,Ny)
for j in 1:Ny
    y[j] = y_min + (j-1)*hy
end
#
Npoints = Nx*Ny
r = zeros(2,Npoints)
ip = 0
idx_ip2xy = zeros(Int64,2,Npoints)
idx_xy2ip = zeros(Int64,Nx,Ny)
for j in 1:Ny
    for i in 1:Nx
        ip = ip + 1
        r[1,ip] = x[i]
        r[2,ip] = y[j]
        idx_ip2xy[1,ip] = i
        idx_ip2xy[2,ip] = j
        idx_xy2ip[i,j] = ip
    end
end
In [54]:
D2x = build_D2_matrix(Nx, hx);
D2y = build_D2_matrix(Ny, hy);
In [57]:
nabla2 = kron(D2x,speye(Ny)) + kron(speye(Nx),D2y);
In [58]:
Matrix(nabla2)
Out[58]:
20×20 Array{Float64,2}:
 -0.5    0.16   0.0    0.0    0.0   …   0.0    0.0    0.0    0.0    0.0 
  0.16  -0.5    0.16   0.0    0.0       0.0    0.0    0.0    0.0    0.0 
  0.0    0.16  -0.5    0.16   0.0       0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.16  -0.5    0.16      0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.0    0.16  -0.5       0.0    0.0    0.0    0.0    0.0 
  0.09   0.0    0.0    0.0    0.0   …   0.0    0.0    0.0    0.0    0.0 
  0.0    0.09   0.0    0.0    0.0       0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.09   0.0    0.0       0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.0    0.09   0.0       0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.0    0.0    0.09      0.0    0.0    0.0    0.0    0.0 
  0.0    0.0    0.0    0.0    0.0   …   0.09   0.0    0.0    0.0    0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.09   0.0    0.0    0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.0    0.09   0.0    0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.0    0.0    0.09   0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.0    0.0    0.0    0.09
  0.0    0.0    0.0    0.0    0.0   …  -0.5    0.16   0.0    0.0    0.0 
  0.0    0.0    0.0    0.0    0.0       0.16  -0.5    0.16   0.0    0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.16  -0.5    0.16   0.0 
  0.0    0.0    0.0    0.0    0.0       0.0    0.0    0.16  -0.5    0.16
  0.0    0.0    0.0    0.0    0.0       0.0    0.0    0.0    0.16  -0.5 

Example of solving Schrodinger equation for harmonic potential

Now, we are ready to solve 2d Schrodinger equation. Like in 1d case, we will convert the differential equation into a matrix eigenvalue problem.

There is one problem that arises, namely the size of Hamiltonian matrix is considerably larger than that of the correnspoding 1d case. We usually don't use the eigh (which uses LAPACK) method in LinearAlgebra method. Even if we can use eigh (assuming we have enough computational resources to use it), we don't usually need all eigenpairs (eigenvalues and eigenvectors) obtained from it. What we need is usually one limited to Nstates (which depend on how many electrons we have) lowest eigenpairs. There are special methods tailored to do such task such as Lanczos, Davidson, conjugate gradient method, etc. In this notebook we will use the so-called LOBPCG (locally-optimal block preconditioned conjugate gradient) method available in the IterativeSolvers package.

In [68]:
Nx = 30
x_min = -5.0
x_max = 5.0
Lx = x_max - x_min
hx = Lx/(Nx-1)
x = zeros(Float64,Nx)
for i in 1:Nx
    x[i] = x_min + (i-1)*hx
end
#
Ny = 30
y_min = -5.0
y_max = 5.0
Ly = y_max - y_min
hy = Ly/(Ny-1)
y = zeros(Float64,Ny)
for j in 1:Ny
    y[j] = y_min + (j-1)*hy
end
#
Npoints = Nx*Ny
r = zeros(2,Npoints)
ip = 0
idx_ip2xy = zeros(Int64,2,Npoints)
idx_xy2ip = zeros(Int64,Nx,Ny)
for j in 1:Ny
    for i in 1:Nx
        ip = ip + 1
        r[1,ip] = x[i]
        r[2,ip] = y[j]
        idx_ip2xy[1,ip] = i
        idx_ip2xy[2,ip] = j
        idx_xy2ip[i,j] = ip
    end
end
In [69]:
Vpot = zeros(Float64,Npoints)
for ip in 1:Npoints
    Vpot[ip] = pot_harmonic(r[1,ip], r[2,ip])
end
In [70]:
D2x = build_D2_matrix(Nx, hx);
D2y = build_D2_matrix(Ny, hy);
nabla2 = kron(D2x,speye(Ny)) + kron(speye(Nx),D2y);
In [71]:
Ham = -0.5*nabla2 + spdiagm( 0 => Vpot );

Find several lowest (say 5) eigenvalues and eigenvectors using LOBPCG method available in IterativeSolvers:

In [72]:
res = lobpcg(Ham, false, 5, maxiter=500)
Out[72]:
Results of LOBPCG Algorithm
 * Algorithm: LOBPCG - CholQR
 * λ: [0.9925118683951702,1.9774185787976424, ...]
 * Residual norm(s): [1.1715886750703023e-5,8.488345014984122e-6, ...]
 * Convergence
   * Iterations: 110
   * Converged: true
   * Iterations limit: 500
In [73]:
res.λ
Out[73]:
5-element Array{Float64,1}:
 0.9925118683951702
 1.9774185787976424
 1.9774185788028917
 2.9469286531790932
 2.9469286535508554

A more modular version: using FD2dGrid struct

A function to build the grid in 1d:

In [74]:
function init_FD1d_grid( X::Tuple{Float64,Float64}, N::Int64 )
    x_min = X[1]
    x_max = X[2]
    L = x_max - x_min
    h = L/(N-1)
    x = zeros(Float64,N)
    for i = 1:N
        x[i] = x_min + (i-1)*h
    end
    return x, h
end
Out[74]:
init_FD1d_grid (generic function with 1 method)

A structure for FD2dGrid:

In [75]:
struct FD2dGrid
    Npoints::Int64
    Nx::Int64
    Ny::Int64
    hx::Float64
    hy::Float64
    x::Array{Float64,1}
    y::Array{Float64,1}
    r::Array{Float64,2}
    idx_ip2xy::Array{Int64,2}
    idx_xy2ip::Array{Int64,2}
end

A constructor of FD2dGrid:

In [76]:
function FD2dGrid(
    x_domain::Tuple{Float64,Float64},
    Nx::Int64,
    y_domain::Tuple{Float64,Float64},
    Ny::Int64
)

    x, hx = init_FD1d_grid(x_domain, Nx)
    y, hy = init_FD1d_grid(y_domain, Ny)
    
    Npoints = Nx*Ny
    r = zeros(2,Npoints)
    ip = 0
    idx_ip2xy = zeros(Int64,2,Npoints)
    idx_xy2ip = zeros(Int64,Nx,Ny)
    for j in 1:Ny
        for i in 1:Nx
            ip = ip + 1
            r[1,ip] = x[i]
            r[2,ip] = y[j]
            idx_ip2xy[1,ip] = i
            idx_ip2xy[2,ip] = j
            idx_xy2ip[i,j] = ip
        end
    end
    return FD2dGrid(Npoints, Nx, Ny, hx, hy, x, y, r, idx_ip2xy, idx_xy2ip)
end
Out[76]:
FD2dGrid

Initialize grid ini two dimensions:

In [77]:
domain_x = (-5.0,5.0)
domain_y = (-5.0,5.0)
Nx = 30
Ny = 30
fdgrid = FD2dGrid( domain_x, Nx, domain_y, Ny );
In [78]:
function build_nabla2_matrix( fdgrid::FD2dGrid )
    Nx = fdgrid.Nx
    hx = fdgrid.hx
    Ny = fdgrid.Ny
    hy = fdgrid.hy
    D2x = build_D2_matrix(Nx, hx)
    D2y = build_D2_matrix(Ny, hy)
    ∇2 = kron(D2x, speye(Ny)) + kron(speye(Nx), D2y)
    return ∇2
end
Out[78]:
build_nabla2_matrix (generic function with 1 method)
In [79]:
∇2 = build_nabla2_matrix(fdgrid);

Return the potential, using linear indexing.

In [81]:
function pot_harmonic(fdgrid::FD2dGrid; ω=1.0)
    Npoints = fdgrid.Npoints
    Vpot = zeros(Npoints)
    for i in 1:Npoints
        x = fdgrid.r[1,i]
        y = fdgrid.r[2,i]
        Vpot[i] = 0.5 * ω^2 *( x^2 + y^2 )
    end
    return Vpot
end
Out[81]:
pot_harmonic (generic function with 2 methods)
In [82]:
Vpot = pot_harmonic(fdgrid);

Build Hamiltonian matrix:

In [83]:
Ham = -0.5*∇2 + spdiagm( 0 => Vpot );

Use LOBPCG to solve the eigenvalue problems.

In [85]:
res = lobpcg(Ham, false, 5)
Out[85]:
Results of LOBPCG Algorithm
 * Algorithm: LOBPCG - CholQR
 * λ: [0.9925118683853325,1.9774185787983265, ...]
 * Residual norm(s): [5.4886365013753944e-6,8.053304187332956e-6, ...]
 * Convergence
   * Iterations: 177
   * Converged: true
   * Iterations limit: 200

The eigenvectors:

In [41]:
X = res.X
Out[41]:
900×5 Array{Float64,2}:
  4.21646e-9   3.76481e-9    2.23903e-8   -1.12358e-8  -2.46906e-8
  2.02217e-9  -2.74825e-10  -2.12011e-8   -7.6763e-9    2.78244e-8
 -1.15772e-9  -3.07111e-9   -3.59698e-11  -4.57389e-9  -2.10054e-8
  1.31404e-9   2.5986e-9    -1.12417e-8   -1.49223e-8  -4.78532e-9
 -2.43817e-9  -2.92505e-9   -8.84855e-9   -4.80738e-8  -2.70507e-8
  3.46374e-9  -3.99319e-8   -2.86019e-8   -1.30818e-7  -5.80764e-8
  7.09718e-9  -8.78363e-8   -2.89855e-8   -3.4649e-7   -2.46276e-7
  3.28269e-8  -2.63439e-7   -5.89776e-8   -7.96083e-7  -6.10753e-7
  8.7944e-8   -5.66837e-7   -1.34565e-7   -1.6486e-6   -1.52977e-6
  1.76609e-7  -1.1389e-6    -1.81786e-7   -3.11597e-6  -3.31514e-6
  3.1461e-7   -2.02104e-6   -1.62784e-7   -5.23736e-6  -6.31015e-6
  5.15044e-7  -3.24276e-6   -4.09984e-8   -7.90928e-6  -1.05552e-5
  7.32594e-7  -4.54134e-6    2.53042e-7   -1.07859e-5  -1.56064e-5
  â‹®                                                               
  5.13886e-7   3.21563e-6    7.17244e-8   -7.92283e-6  -1.05587e-5
  3.25229e-7   2.01936e-6    1.93874e-7   -5.27016e-6  -6.27642e-6
  1.81359e-7   1.14672e-6    1.86536e-7   -3.13309e-6  -3.29141e-6
  8.02125e-8   5.74888e-7    1.27347e-7   -1.68621e-6  -1.53141e-6
  4.07928e-8   2.58934e-7    6.51389e-8   -8.01795e-7  -6.14757e-7
  1.11189e-8   1.01494e-7    2.47018e-8   -3.28834e-7  -2.35988e-7
  1.30632e-8   4.05789e-8    1.62707e-8   -1.28833e-7  -7.20877e-8
 -2.06493e-9   8.35929e-9    9.52268e-9   -2.57268e-8  -1.31719e-8
 -1.44139e-9   9.52057e-9    1.61037e-8   -1.45482e-8  -2.64622e-8
 -3.09455e-9   4.96018e-9   -2.81852e-8    7.52707e-9   2.17467e-8
 -6.5494e-9   -1.92994e-9    4.25464e-8   -1.74379e-8  -2.94917e-8
  1.54843e-9  -5.96422e-9   -3.01934e-8   -4.25602e-9   5.23181e-9

This is the eigenvalues:

In [42]:
res.λ
Out[42]:
5-element Array{Float64,1}:
 0.992511868385896 
 1.9774185787949063
 1.9774185788034842
 2.946928653227466 
 2.946928654618251 
$n_{x}$ $n_{y}$ $E_{n_{x},n_{y}}$
0 0 1.0
1 0 2.0
0 1 2.0
1 1 3.0
2 0 3.0
0 2 3.0
2 1 4.0

Visualization of the wave functions (orbitals)

Need to use reshape.

1st orbital

In [58]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(fdgrid.x, fdgrid.y, reshape(res.X[:,1], (fdgrid.Nx, fdgrid.Ny)) )
  )
)
Out[58]:

2nd orbital

In [57]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(fdgrid.x, fdgrid.y, reshape(res.X[:,2], (fdgrid.Nx, fdgrid.Ny)) )
  )
)
Out[57]:

3rd orbital

In [56]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(fdgrid.x, fdgrid.y, reshape(res.X[:,3], (fdgrid.Nx, fdgrid.Ny)) )
  )
)
Out[56]:

4th orbital

In [86]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(fdgrid.x, fdgrid.y, reshape(res.X[:,4], (fdgrid.Nx, fdgrid.Ny)) )
  )
)
Out[86]:

5th orbital

In [55]:
@pgf Axis({ height = "10cm", width = "10cm", view=(20,10), colorbar },
  Plot3(
    {
        surf,
    },
    Coordinates(fdgrid.x, fdgrid.y, reshape(res.X[:,5], (fdgrid.Nx, fdgrid.Ny)) )
  )
)
Out[55]:
In [ ]:

In [ ]:

In [ ]: